import sys, os
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
import json
from preprossData.utils.crop_padding import crop_resize_save

def gather_data(data_path, tar_path, prefix='szEyeROP'):
    """
    对新数据集进行预处理：
      - 数据集根目录下存在若干类别文件夹，每个文件夹名称与 folder2text 中的 key 对应，
        文件夹内所有图像的标签均为 folder2text 中对应的诊断文本；
      - 将所有处理后的图像统一保存到 tar_path/images 下，
        新图像名称格式为 {prefix}_{原始文件名}.png（保留原始文件名，仅修改后缀并加上前缀）；
      - 生成标注信息字典，并将其保存到 tar_path/annotations.json。

    参数：
        data_path (str): 新数据集根目录，要求其下包含各类别文件夹。
        tar_path (str): 预处理后数据存放目录。
        prefix (str): 处理后图片名称的前缀，默认为 "szEyeROP".

    返回：
        dict: 标注信息字典，键为新图像名称，值为图像的相关信息。
    """
    os.makedirs(tar_path, exist_ok=True)
    images_dir = os.path.join(tar_path, 'images')
    os.makedirs(images_dir, exist_ok=True)

    data_dict = {}

    folder2text = {
        'laser scars': 'laser scars',
        'Normal': "normal",
        "Stage1": "Retinopathy of Prematurity stage 1",
        "Stage2": "Retinopathy of Prematurity stage 2",
        "Stage3": "Retinopathy of Prematurity stage 3"
    }

    # 遍历 folder2text 中的每个类别文件夹
    for class_folder, text_label in folder2text.items():
        image_folder = os.path.join(data_path,'image', class_folder)
        if not os.path.exists(image_folder):
            raise ValueError(f"Folder {image_folder} does not exist.")
        # 遍历该类别文件夹中的所有图像文件
        for image_file in sorted(os.listdir(image_folder)):
            if not image_file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.ppm')):
                continue
            src_image_path = os.path.join(image_folder, image_file)
            # 保留原始文件名（不含扩展名），仅加上前缀并统一保存为 PNG 格式
            base_name = os.path.splitext(image_file)[0]
            new_image_name = f"{prefix}_{base_name}.png"
            dest_image_path = os.path.join(images_dir, new_image_name)
            
            # 对图像进行裁剪与 resize，并保存为 PNG 格式
            crop_info=crop_resize_save(
                image_path=src_image_path,
                save_path=dest_image_path,
                resize=(224, 224),
                crop_threshold=25
            )
            
            # 保存标注信息（这里 image_path 为相对于 tar_path 的路径）
            data_dict[new_image_name] = {
                'image_name': new_image_name,
                'image_path': os.path.join('images', new_image_name),
                'original_path': os.path.relpath(src_image_path, data_path),
                
                'crop_info':crop_info,
                'diagnosis': {
                    'text': text_label
                }
            }
    
    # 保存标注信息到 tar_path/annotations.json
    annotations_path = os.path.join(tar_path, 'annotations.json')
    with open(annotations_path, 'w', encoding='utf-8') as f:
        json.dump(data_dict, f, indent=4)
    
    return data_dict

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="新数据集预处理")
    parser.add_argument("--data_path", type=str, default="/mnt/e/Document_Workspace/Dataset/public_dataset/shenZhenEyeROP/",
                        help="新数据集根目录，要求其下包含类别文件夹")
    parser.add_argument("--tar_path", type=str, default="/mnt/e/Document_Workspace/Dataset/processed224_new/shenZhenEyeROP",
                        help="预处理后数据存放目录")
    parser.add_argument("--prefix", type=str, default="szEyeROP",
                        help="处理后图片名称的前缀，默认为 'szEyeROP'")
    args = parser.parse_args()
    
    annotations = gather_data(args.data_path, args.tar_path, prefix=args.prefix)
    print("Preprocessing completed. Annotations saved.")
